import os
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def convert_pytorch_to_onnx(model_path, output_dir=None):
    """
    将PyTorch模型转换为ONNX格式
    
    Args:
        model_path: PyTorch模型路径
        output_dir: 输出目录，默认为model_path/onnx
    """
    if output_dir is None:
        output_dir = os.path.join(model_path, "onnx")
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    onnx_model_path = os.path.join(output_dir, "model.onnx")
    
    logger.info(f"加载PyTorch模型: {model_path}")
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
    model = AutoModelForTokenClassification.from_pretrained(model_path, local_files_only=True)
    
    # 将模型设置为评估模式
    model.eval()
    
    # 创建示例输入
    dummy_input_ids = torch.ones(1, 128, dtype=torch.long)
    dummy_attention_mask = torch.ones(1, 128, dtype=torch.long)
    dummy_token_type_ids = torch.zeros(1, 128, dtype=torch.long)
    
    # 准备输入名称和动态轴
    input_names = ["input_ids", "attention_mask", "token_type_ids"]
    output_names = ["logits"]
    dynamic_axes = {
        "input_ids": {0: "batch_size", 1: "sequence"},
        "attention_mask": {0: "batch_size", 1: "sequence"},
        "token_type_ids": {0: "batch_size", 1: "sequence"},
        "logits": {0: "batch_size", 1: "sequence"}
    }
    
    logger.info(f"转换模型为ONNX格式: {onnx_model_path}")
    torch.onnx.export(
        model,
        (dummy_input_ids, dummy_attention_mask, dummy_token_type_ids),
        onnx_model_path,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        opset_version=12,
        do_constant_folding=True
    )
    
    logger.info(f"ONNX模型已保存到: {onnx_model_path}")
    return onnx_model_path

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("用法: python convert_pytorch_to_onnx.py <模型路径> [输出目录]")
        sys.exit(1)
    
    model_path = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else None
    
    try:
        convert_pytorch_to_onnx(model_path, output_dir)
        print(f"转换成功!")
    except Exception as e:
        logger.error(f"转换失败: {str(e)}")
        print(f"错误: {str(e)}")
        sys.exit(1) 